from __future__ import print_function

from keras.models import Sequential
from keras.layers.convolutional import Convolution2D, MaxPooling2D, ZeroPadding2D
from keras.layers.core import Activation, Dense, Flatten, Dropout, Reshape
from keras.regularizers import l2
from keras.layers.recurrent import LSTM
from keras.layers.advanced_activations import LeakyReLU, ELU, PReLU
from keras.layers.normalization import BatchNormalization
from keras import backend as K

from keras_extras.extra import TimeDistributedFlatten, TimeDistributedMaxPooling2D, TimeDistributedConvolution2D


def scale(x):
    return (x - K.mean(x)) / K.std(x)
    # return x


def get_model():
    conv = Sequential()
    conv.add(Activation(activation=scale, input_shape=(30, 1, 128, 128)))
    conv.add(TimeDistributedConvolution2D(64, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    # model.add(TimeDistributedConvolution2D(64, 3, 3))
    # model.add(Activation('relu'))
    conv.add(TimeDistributedMaxPooling2D(pool_size=(2, 2)))
    conv.add(TimeDistributedConvolution2D(128, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    # model.add(TimeDistributedConvolution2D(128, 3, 3))
    # model.add(Activation('relu'))
    conv.add(TimeDistributedMaxPooling2D(pool_size=(2, 2)))
    conv.add(TimeDistributedConvolution2D(256, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    # model.add(TimeDistributedConvolution2D(256, 3, 3))
    # model.add(Activation('relu'))
    conv.add(TimeDistributedMaxPooling2D(pool_size=(2, 2)))
    conv.add(TimeDistributedConvolution2D(512, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    # model.add(TimeDistributedConvolution2D(512, 3, 3))
    # model.add(Activation('relu'))
    conv.add(TimeDistributedMaxPooling2D(pool_size=(2, 2)))
    conv.add(Activation('relu'))
    conv.add(TimeDistributedConvolution2D(512, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    # model.add(TimeDistributedConvolution2D(512, 3, 3))
    # model.add(Activation('relu'))
    conv.add(TimeDistributedMaxPooling2D(pool_size=(2, 2)))
    conv.add(Activation('relu'))
    conv.add(TimeDistributedConvolution2D(512, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    # model.add(TimeDistributedConvolution2D(512, 3, 3))
    # model.add(Activation('relu'))
    conv.add(TimeDistributedMaxPooling2D(pool_size=(2, 2)))
    conv.add(Activation('relu'))
    # model.add(TimeDistributedConvolution2D(512, 3, 3, border_mode='same'))
    # model.add(Activation('relu'))
    # # model.add(TimeDistributedConvolution2D(512, 3, 3))
    # # model.add(Activation('relu'))
    # model.add(TimeDistributedMaxPooling2D(pool_size=(2, 2)))
    # model.add(Activation('relu'))
    # model.add(TimeDistributedConvolution2D(512, 3, 3, border_mode='same'))
    # model.add(Activation('relu'))
    # # model.add(TimeDistributedConvolution2D(512, 3, 3))
    # # model.add(Activation('relu'))
    # model.add(TimeDistributedMaxPooling2D(pool_size=(2, 2)))
    # model.add(Activation('relu'))
    conv.add(TimeDistributedFlatten())
    conv.add(LSTM(512, return_sequences=True))
    conv.add(Dropout(0.5))
    conv.add(Dense(1))
    return conv
